from vision.core.image_reader import TrainImageReader
import datetime
import os
from vision.core.models import PNet,RNet,ONet,LossFn
import torch
from torch.autograd import Variable
import vision.core.image_tools as image_tools
import time
import numpy as np
from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR,StepLR

def compute_accuracy(prob_cls, gt_cls):
    prob_cls = torch.squeeze(prob_cls)
    gt_cls = torch.squeeze(gt_cls)

    #we only need the detection which >= 0
    mask = torch.ge(gt_cls,0)
    #get valid element
    valid_gt_cls = torch.masked_select(gt_cls,mask)
    valid_prob_cls = torch.masked_select(prob_cls,mask)
    size = min(valid_gt_cls.size()[0], valid_prob_cls.size()[0])
    prob_ones = torch.ge(valid_prob_cls,0.6).float()
    right_ones = torch.eq(prob_ones,valid_gt_cls).float()

    return torch.div(torch.mul(torch.sum(right_ones),float(1.0)),float(size))


def train_pnet(model_store_path,resume, end_epoch, train_loader, test_loader,
              frequent=50,base_lr=0.01,use_cuda=True):

    if not os.path.exists(model_store_path):
        os.makedirs(model_store_path)

    lossfn = LossFn()
    net = PNet(is_train=True, use_cuda=use_cuda)
    if resume:
        net.load_state_dict(torch.load(resume))
    if use_cuda:
        net.cuda()

    optimizer = torch.optim.Adam(net.parameters(), lr=base_lr)


    for cur_epoch in range(1,end_epoch+1):
        accuracy_list=[]
        cls_loss_list=[]
        bbox_loss_list=[]
        # landmark_loss_list=[]

        print("SAVE TEST")
        torch.save(net.state_dict(), os.path.join(model_store_path,"pnet_epoch_%d.pt" % cur_epoch))
        torch.save(net, os.path.join(model_store_path,"pnet_epoch_model_%d.pkl" % cur_epoch))

        net.train()
        back_time = time.time()
        for batch_idx, data in enumerate(train_loader):
            start_time = time.time()
            loader_time = start_time-back_time
            image, gt_label, gt_bbox, gt_landmark = data

            im_tensor = Variable(image)
            gt_label = Variable(gt_label)

            gt_bbox = Variable(gt_bbox)
            # gt_landmark = Variable(torch.from_numpy(gt_landmark).float())

            if use_cuda:
                im_tensor = im_tensor.cuda()
                gt_label = gt_label.cuda()
                gt_bbox = gt_bbox.cuda()
                # gt_landmark = gt_landmark.cuda()
            cvt_time = time.time()
            cls_pred, box_offset_pred = net(im_tensor)
            forward_time = time.time()
            # all_loss, cls_loss, offset_loss = lossfn.loss(gt_label=label_y,gt_offset=bbox_y, pred_label=cls_pred, pred_offset=box_offset_pred)

            cls_loss = lossfn.cls_loss(gt_label,cls_pred)
            box_offset_loss = lossfn.box_loss(gt_label,gt_bbox,box_offset_pred)
            # landmark_loss = lossfn.landmark_loss(gt_label,gt_landmark,landmark_offset_pred)

            all_loss = cls_loss*1.0+box_offset_loss*0.5

            if batch_idx % frequent == 0:
                accuracy=compute_accuracy(cls_pred,gt_label)

                show1 = accuracy.item()
                show2 = cls_loss.item()
                show3 = box_offset_loss.item()
                show5 = all_loss.item()

                print("TRAIN %s : Epoch: %d, Step: %d, accuracy: %s, det loss: %s, bbox loss: %s, all_loss: %s, lr:%s "%(datetime.datetime.now(),cur_epoch,batch_idx, show1,show2,show3,show5,base_lr))

            optimizer.zero_grad()
            all_loss.backward()
            optimizer.step()
            back_time = time.time()
            print("loader:%f,cvt:%f,forward:%f,backward:%f"%(loader_time, cvt_time - start_time,forward_time - cvt_time,back_time-forward_time))
        net.eval()
        for batch_idx, data in enumerate(test_loader):
            image, gt_label, gt_bbox, gt_landmark = data

            im_tensor = Variable(image)
            gt_label = Variable(gt_label)
            gt_bbox = Variable(gt_bbox)
            # gt_landmark = Variable(torch.from_numpy(gt_landmark).float())

            if use_cuda:
                im_tensor = im_tensor.cuda()
                gt_label = gt_label.cuda()
                gt_bbox = gt_bbox.cuda()
                # gt_landmark = gt_landmark.cuda()

            cls_pred, box_offset_pred = net(im_tensor)
            # all_loss, cls_loss, offset_loss = lossfn.loss(gt_label=label_y,gt_offset=bbox_y, pred_label=cls_pred, pred_offset=box_offset_pred)

            cls_loss = lossfn.cls_loss(gt_label,cls_pred)
            box_offset_loss = lossfn.box_loss(gt_label,gt_bbox,box_offset_pred)
            # landmark_loss = lossfn.landmark_loss(gt_label,gt_landmark,landmark_offset_pred)

            accuracy = compute_accuracy(cls_pred,gt_label)
            accuracy_list.append(accuracy.item())
            cls_loss_list.append(cls_loss.item())
            bbox_loss_list.append(box_offset_loss.item())

        accuracy_avg = np.mean(np.array(accuracy_list))
        cls_loss_avg = np.mean(np.array(cls_loss_list))
        bbox_loss_avg = np.mean(np.array(bbox_loss_list))
        # landmark_loss_avg = torch.mean(torch.cat(landmark_loss_list))

        show6 = accuracy_avg
        show7 = cls_loss_avg
        show8 = bbox_loss_avg

        print("TEST %s Epoch: %d, accuracy: %s, cls loss: %s, bbox loss: %s" % (datetime.datetime.now(), cur_epoch, show6, show7, show8))

        torch.save(net.state_dict(), os.path.join(model_store_path,"pnet_epoch_%d_%.6f.pt" % (cur_epoch,accuracy_avg)))
        torch.save(net, os.path.join(model_store_path,"pnet_epoch_model_%d_%.6f.pkl" % (cur_epoch,accuracy_avg)))




def train_rnet(model_store_path, resume, end_epoch, train_loader, test_loader,
              frequent=50,base_lr=0.01,use_cuda=True):

    if not os.path.exists(model_store_path):
        os.makedirs(model_store_path)

    lossfn = LossFn()
    net = RNet(is_train=True, use_cuda=use_cuda)
    if resume:
        net.load_state_dict(torch.load(resume))
    if use_cuda:
        net.cuda()

    optimizer = torch.optim.Adam(net.parameters(), lr=base_lr)


    for cur_epoch in range(1,end_epoch+1):
        accuracy_list=[]
        cls_loss_list=[]
        bbox_loss_list=[]
        # landmark_loss_list=[]

        print("SAVE TEST")
        torch.save(net.state_dict(), os.path.join(model_store_path,"rnet_epoch_%d.pt" % cur_epoch))
        torch.save(net, os.path.join(model_store_path,"rnet_epoch_model_%d.pkl" % cur_epoch))


        net.eval()
        for batch_idx, data in enumerate(test_loader):
            image, gt_label, gt_bbox, gt_landmark = data

            im_tensor = Variable(image)
            gt_label = Variable(gt_label)

            gt_bbox = Variable(gt_bbox)
            # gt_landmark = Variable(torch.from_numpy(gt_landmark).float())

            if use_cuda:
                im_tensor = im_tensor.cuda()
                gt_label = gt_label.cuda()
                gt_bbox = gt_bbox.cuda()
                # gt_landmark = gt_landmark.cuda()

            cls_pred, box_offset_pred = net(im_tensor)
            # all_loss, cls_loss, offset_loss = lossfn.loss(gt_label=label_y,gt_offset=bbox_y, pred_label=cls_pred, pred_offset=box_offset_pred)

            cls_loss = lossfn.cls_loss(gt_label, cls_pred)
            box_offset_loss = lossfn.box_loss(gt_label, gt_bbox, box_offset_pred)
            accuracy = compute_accuracy(cls_pred,gt_label)
            accuracy_list.append(accuracy.item())
            cls_loss_list.append(cls_loss.item())
            bbox_loss_list.append(box_offset_loss.item())

        accuracy_avg = np.mean(np.array(accuracy_list))
        cls_loss_avg = np.mean(np.array(cls_loss_list))
        bbox_loss_avg = np.mean(np.array(bbox_loss_list))
        # landmark_loss_avg = torch.mean(torch.cat(landmark_loss_list))

        show6 = accuracy_avg
        show7 = cls_loss_avg
        show8 = bbox_loss_avg

        print("TEST %s Epoch: %d, accuracy: %s, cls loss: %s, bbox loss: %s" % (datetime.datetime.now(), cur_epoch, show6, show7, show8))

        torch.save(net.state_dict(), os.path.join(model_store_path,"rnet_epoch_%d_%.6f.pt" % (cur_epoch,accuracy_avg)))
        torch.save(net, os.path.join(model_store_path,"rnet_epoch_model_%d_%.6f.pkl" % (cur_epoch,accuracy_avg)))

        net.train()
        for batch_idx, data in enumerate(train_loader):
            image, gt_label, gt_bbox, gt_landmark = data

            im_tensor = Variable(image)
            gt_label = Variable(gt_label)

            gt_bbox = Variable(gt_bbox)
            # gt_landmark = Variable(torch.from_numpy(gt_landmark).float())

            if use_cuda:
                im_tensor = im_tensor.cuda()
                gt_label = gt_label.cuda()
                gt_bbox = gt_bbox.cuda()
                # gt_landmark = gt_landmark.cuda()

            cls_pred, box_offset_pred = net(im_tensor)
            # all_loss, cls_loss, offset_loss = lossfn.loss(gt_label=label_y,gt_offset=bbox_y, pred_label=cls_pred, pred_offset=box_offset_pred)

            cls_loss = lossfn.cls_loss(gt_label,cls_pred)
            box_offset_loss = lossfn.box_loss(gt_label,gt_bbox,box_offset_pred)
            # landmark_loss = lossfn.landmark_loss(gt_label,gt_landmark,landmark_offset_pred)

            all_loss = cls_loss*1.0+box_offset_loss*10

            if batch_idx % frequent == 0:
                accuracy=compute_accuracy(cls_pred,gt_label)

                show1 = accuracy.item()
                show2 = cls_loss.item()
                show3 = box_offset_loss.item()
                show5 = all_loss.item()

                print("TRAIN %s : Epoch: %d, Step: %d, accuracy: %s, det loss: %s, bbox loss: %s, all_loss: %s, lr:%s "%(datetime.datetime.now(),cur_epoch,batch_idx, show1,show2,show3,show5,base_lr))

            optimizer.zero_grad()
            all_loss.backward()
            optimizer.step()


def train_onet(model_store_path,resume, end_epoch, train_loader, test_loader,
              frequent=50,base_lr=0.01,use_cuda=True):

    if not os.path.exists(model_store_path):
        os.makedirs(model_store_path)

    lossfn = LossFn()
    net = ONet(is_train=True, use_cuda=use_cuda)

    if resume:
        net.load_state_dict(torch.load(resume))
    if use_cuda:
        net.cuda()
    last_epoch = -1
    optimizer = torch.optim.Adam(net.parameters(), lr=base_lr)
    # optimizer = torch.optim.SGD(net.parameters(), lr=base_lr,momentum=0.9)
    # lr_scheduler = CosineAnnealingLR(optimizer, T_max=5, eta_min=base_lr * 0.1,last_epoch=last_epoch)
    lr_scheduler = StepLR(optimizer,step_size=10,gamma=0.1)
    # optimizer = torch.optim.RMSprop(net.parameters())

    for cur_epoch in range(last_epoch + 1,end_epoch):
        lr_scheduler.step()
        accuracy_list=[]
        cls_loss_list=[]
        bbox_loss_list=[]
        landmark_loss_list=[]

        print("SAVE TEST")
        torch.save(net.state_dict(), os.path.join(model_store_path,"onet_epoch_%d.pt" % cur_epoch))
        torch.save(net, os.path.join(model_store_path,"onet_epoch_model_%d.pkl" % cur_epoch))

        net.train()
        back_time = time.time()
        for batch_idx, data in enumerate(train_loader):

            start_time = time.time()
            loader_time = start_time - back_time
            image, gt_label, gt_bbox, gt_landmark = data
            im_tensor = Variable(image)
            gt_label = Variable(gt_label)

            gt_bbox = Variable(gt_bbox)
            gt_landmark = Variable(gt_landmark)

            if use_cuda:

                im_tensor = im_tensor.cuda()
                gt_label = gt_label.cuda()
                gt_bbox = gt_bbox.cuda()
                gt_landmark = gt_landmark.cuda()

            cvt_time = time.time()

            cls_pred, box_offset_pred, landmark_offset_pred = net(im_tensor)
            forward_time = time.time()
            # all_loss, cls_loss, offset_loss = lossfn.loss(gt_label=label_y,gt_offset=bbox_y, pred_label=cls_pred, pred_offset=box_offset_pred)

            cls_loss = lossfn.cls_loss(gt_label,cls_pred)
            box_offset_loss = lossfn.box_loss(gt_label,gt_bbox,box_offset_pred)
            landmark_loss = lossfn.landmark_loss(gt_label,gt_landmark,landmark_offset_pred)

            all_loss = cls_loss+box_offset_loss*10+landmark_loss*1000
            # all_loss = landmark_loss * 100
            if batch_idx % frequent == 0:
                accuracy=compute_accuracy(cls_pred,gt_label)

                show1 = accuracy.item()
                # show1 = 0
                show2 = cls_loss.item()
                show3 = box_offset_loss.item()
                show4 = landmark_loss.item()
                show5 = all_loss.item()
                for param_group in optimizer.param_groups:
                    cur_lr=param_group['lr']
                print("TRAIN %s : Epoch: %d, Step: %d, accuracy: %.6f, det loss: %.6f, bbox loss: %.6f, landmark loss: %.6f, all_loss: %.6f, lr:%.6f "%
                      (datetime.datetime.now(),cur_epoch,batch_idx, show1,show2,show3,show4,show5,cur_lr))

            optimizer.zero_grad()
            all_loss.backward()
            optimizer.step()

            back_time = time.time()
            # print("loader:%f,cvt:%f,forward:%f,backward:%f"%(loader_time, cvt_time - start_time,forward_time - cvt_time,back_time-forward_time))

        net.eval()
        for batch_idx, data in enumerate(test_loader):
            image, gt_label, gt_bbox, gt_landmark = data

            im_tensor = Variable(image)
            gt_label = Variable(gt_label)
            gt_bbox = Variable(gt_bbox)
            # gt_landmark = Variable(torch.from_numpy(gt_landmark).float())
            gt_landmark = Variable(gt_landmark)
            if use_cuda:
                im_tensor = im_tensor.cuda()
                gt_label = gt_label.cuda()
                gt_bbox = gt_bbox.cuda()
                gt_landmark = gt_landmark.cuda()

            cls_pred, box_offset_pred, landmark_offset_pred = net(im_tensor)

            cls_loss = lossfn.cls_loss(gt_label,cls_pred)
            box_offset_loss = lossfn.box_loss(gt_label,gt_bbox,box_offset_pred)
            landmark_loss = lossfn.landmark_loss(gt_label,gt_landmark,landmark_offset_pred)

            accuracy = compute_accuracy(cls_pred,gt_label)
            accuracy_list.append(accuracy.item())
            cls_loss_list.append(cls_loss.item())
            bbox_loss_list.append(box_offset_loss.item())
            landmark_loss_list.append(landmark_loss.item())

        accuracy_avg = np.mean(np.array(accuracy_list))
        cls_loss_avg = np.mean(np.array(cls_loss_list))
        bbox_loss_avg = np.mean(np.array(bbox_loss_list))
        landmark_loss_avg = np.mean(np.array(landmark_loss_list))

        show6 = accuracy_avg
        # show6 = 0
        show7 = cls_loss_avg
        show8 = bbox_loss_avg
        show9 = landmark_loss_avg

        print("TEST %s Epoch: %d, accuracy: %s, cls loss: %s, bbox loss: %s, ldmark loss: %s" % (datetime.datetime.now(), cur_epoch, show6, show7, show8, show9))
        torch.save(net.state_dict(), os.path.join(model_store_path,"onet_epoch_%d_%.6f_%.4f.pt" % (cur_epoch,accuracy_avg,landmark_loss_avg)))
        torch.save(net, os.path.join(model_store_path,"onet_epoch_model_%d_%.6f_%.4f.pkl" % (cur_epoch,accuracy_avg,landmark_loss_avg)))

